from collections.abc import Iterator
from contextlib import contextmanager

from myapp.config import SQLALCHEMY_DATABASE_URI
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import NullPool


@contextmanager
def session_scope() -> Iterator[Session]:
    """Provide a transactional scope around a series of operations."""
    database_uri = SQLALCHEMY_DATABASE_URI
    engine = create_engine(database_uri, poolclass=NullPool)
    session_class = sessionmaker()
    session_class.configure(bind=engine)
    session = session_class()

    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def scan_table(model, filters=None, page_size=100):
    """分页查询遍历数据库"""
    last_id = 0
    # 分页查询遍历数据库
    with session_scope() as dbsession:
        while True:
            q = dbsession.query(model).filter(model.id > last_id)

            if filters is not None:
                if isinstance(filters, list):
                    filters = tuple(filters)
                q = q.filter(*filters) if isinstance(filters, tuple) else q.filter(filters)

            results = q.order_by('id').limit(page_size).all()
            if results is None or len(results) == 0:
                break

            yield results, dbsession

            if len(results) < page_size:
                break
            last_id = results[-1].id
